{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b10f35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deep Research Agent\n",
    "## 1. Clarifying questions when the query is vague\n",
    "## 2. Multiple parallel searches\n",
    "## 3. Report quality evaluation\n",
    "## 4. Automatic retry if evaluation fails\n",
    "## 5. Email delivery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04f30ca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from agents import Agent, WebSearchTool, trace, Runner, function_tool\n",
    "from agents.model_settings import ModelSettings\n",
    "from pydantic import BaseModel, Field\n",
    "from dotenv import load_dotenv\n",
    "import asyncio\n",
    "import requests\n",
    "import os\n",
    "from typing import Dict\n",
    "from IPython.display import display, Markdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e36b4485",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d7de67a",
   "metadata": {},
   "source": [
    "## Clarifier Agent\n",
    "\n",
    "Asks clarifying questions when the query is vague or broad."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d23a26b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClarificationOutput(BaseModel):\n",
    "    needs_clarification: bool = Field(description=\"Whether the query needs clarification\")\n",
    "    questions: list[str] = Field(description=\"List of clarifying questions to ask the user\")\n",
    "\n",
    "CLARIFIER_INSTRUCTIONS = \"\"\"You are a research clarification agent. Given a research query, decide if it needs clarification.\n",
    "\n",
    "A query needs clarification if it is:\n",
    "- Too broad or vague\n",
    "- Missing key context like timeframe, location, or specific focus\n",
    "- Ambiguous in intent\n",
    "\n",
    "If clarification is needed, generate 2-3 specific questions that would help narrow the scope and improve research quality.\n",
    "If the query is clear and specific enough, set needs_clarification to false and return empty questions list.\"\"\"\n",
    "\n",
    "clarifier_agent = Agent(\n",
    "    name=\"ClarifierAgent\",\n",
    "    instructions=CLARIFIER_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ClarificationOutput,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6de38c2",
   "metadata": {},
   "source": [
    "## Search Planner Agent\n",
    "\n",
    "Creates a strategic plan of web searches to perform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ef3bfe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "HOW_MANY_SEARCHES = 5\n",
    "\n",
    "class WebSearchItem(BaseModel):\n",
    "    reason: str = Field(description=\"Why this search is important\")\n",
    "    query: str = Field(description=\"The search term to use\")\n",
    "\n",
    "class WebSearchPlan(BaseModel):\n",
    "    searches: list[WebSearchItem] = Field(description=\"List of web searches to perform\")\n",
    "\n",
    "PLANNER_INSTRUCTIONS = f\"\"\"You are a research planning agent. Given a query, create {HOW_MANY_SEARCHES} strategic web searches to best answer it.\n",
    "\n",
    "Make searches specific and complementary, covering different angles of the topic.\"\"\"\n",
    "\n",
    "planner_agent = Agent(\n",
    "    name=\"PlannerAgent\",\n",
    "    instructions=PLANNER_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=WebSearchPlan,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0277d19",
   "metadata": {},
   "source": [
    "## Search Agent\n",
    "\n",
    "Performs web searches and summarizes results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4451a6a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEARCH_INSTRUCTIONS = \"\"\"You are a research assistant. Given a search term, search the web and produce a concise summary.\n",
    "\n",
    "Keep summaries to 2-3 paragraphs and under 300 words. Capture main points succinctly.\"\"\"\n",
    "\n",
    "search_agent = Agent(\n",
    "    name=\"SearchAgent\",\n",
    "    instructions=SEARCH_INSTRUCTIONS,\n",
    "    tools=[WebSearchTool(search_context_size=\"low\")],\n",
    "    model=\"gpt-4o-mini\",\n",
    "    model_settings=ModelSettings(tool_choice=\"required\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9816223",
   "metadata": {},
   "source": [
    "## Report Writer Agent\n",
    "\n",
    "Synthesizes search results into a comprehensive report."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "549dfe30",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReportData(BaseModel):\n",
    "    short_summary: str = Field(description=\"2-3 sentence summary of findings\")\n",
    "    markdown_report: str = Field(description=\"The full report in markdown\")\n",
    "    follow_up_questions: list[str] = Field(description=\"Suggested topics to research further\")\n",
    "\n",
    "WRITER_INSTRUCTIONS = \"\"\"You are a senior researcher writing a comprehensive report.\n",
    "\n",
    "You will receive the original query and research summaries. Create a detailed report that:\n",
    "- Provides clear structure and flow\n",
    "- Synthesizes all findings cohesively\n",
    "- Uses markdown formatting\n",
    "- Is lengthy and detailed, aiming for 1000+ words\n",
    "- Includes an executive summary, main findings, and conclusions\"\"\"\n",
    "\n",
    "writer_agent = Agent(\n",
    "    name=\"WriterAgent\",\n",
    "    instructions=WRITER_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ReportData,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "965b920c",
   "metadata": {},
   "source": [
    "## Report Evaluator Agent\n",
    "\n",
    "Evaluates if the report meets quality standards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b895d1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EvaluationResult(BaseModel):\n",
    "    is_acceptable: bool = Field(description=\"Whether the report meets quality standards\")\n",
    "    feedback: str = Field(description=\"Specific feedback on what needs improvement\")\n",
    "\n",
    "EVALUATOR_INSTRUCTIONS = \"\"\"You are a report quality evaluator. Assess if the research report is:\n",
    "- Comprehensive and well-structured\n",
    "- Clearly written and easy to follow\n",
    "- Sufficiently detailed (at least 800 words)\n",
    "- Properly addresses the original query\n",
    "- Contains actionable insights\n",
    "\n",
    "Provide specific feedback on what needs improvement if the report is not acceptable.\"\"\"\n",
    "\n",
    "evaluator_agent = Agent(\n",
    "    name=\"EvaluatorAgent\",\n",
    "    instructions=EVALUATOR_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=EvaluationResult,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fb52e9b",
   "metadata": {},
   "source": [
    "## Email Agent Agent\n",
    "\n",
    "Formats and sends the final report via email."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae22a7dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "@function_tool\n",
    "def send_email(subject: str, html_body: str) -> Dict[str, str]:\n",
    "    mailgun_domain = os.getenv(\"MAILGUN_DOMAIN\")\n",
    "    mailgun_api_key = os.getenv(\"MAILGUN_API_KEY\")\n",
    "    from_email = os.getenv(\"MAILGUN_FROM_EMAIL\", \"noreply@\" + mailgun_domain)\n",
    "    to_email = os.getenv(\"MAILGUN_TO_EMAIL\")\n",
    "    \n",
    "    response = requests.post(\n",
    "        f\"https://api.mailgun.net/v3/{mailgun_domain}/messages\",\n",
    "        auth=(\"api\", mailgun_api_key),\n",
    "        data={\n",
    "            \"from\": from_email,\n",
    "            \"to\": to_email,\n",
    "            \"subject\": subject,\n",
    "            \"html\": html_body\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    return {\"status\": \"success\" if response.status_code == 200 else \"failed\"}\n",
    "\n",
    "EMAIL_INSTRUCTIONS = \"\"\"You are an email formatter. Convert the research report into a well-formatted HTML email.\n",
    "\n",
    "Create a clean, professional HTML layout with appropriate subject line.\"\"\"\n",
    "\n",
    "email_agent = Agent(\n",
    "    name=\"EmailAgent\",\n",
    "    instructions=EMAIL_INSTRUCTIONS,\n",
    "    tools=[send_email],\n",
    "    model=\"gpt-4o-mini\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "567de745",
   "metadata": {},
   "source": [
    "## Research Workflow Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db47db8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def clarify_query(query: str):\n",
    "    print(\"Checking if query needs clarification...\")\n",
    "    result = await Runner.run(clarifier_agent, f\"Query: {query}\")\n",
    "    return result.final_output\n",
    "\n",
    "async def plan_searches(query: str):\n",
    "    print(\"Planning searches...\")\n",
    "    result = await Runner.run(planner_agent, f\"Query: {query}\")\n",
    "    print(f\"Will perform {len(result.final_output.searches)} searches\")\n",
    "    return result.final_output\n",
    "\n",
    "async def perform_searches(search_plan: WebSearchPlan):\n",
    "    print(\"Searching...\")\n",
    "    tasks = [asyncio.create_task(search(item)) for item in search_plan.searches]\n",
    "    results = await asyncio.gather(*tasks)\n",
    "    print(\"Finished searching\")\n",
    "    return results\n",
    "\n",
    "async def search(item: WebSearchItem):\n",
    "    input = f\"Search term: {item.query}\\nReason: {item.reason}\"\n",
    "    result = await Runner.run(search_agent, input)\n",
    "    return result.final_output\n",
    "\n",
    "async def write_report(query: str, search_results: list[str], previous_feedback: str = \"\"):\n",
    "    print(\"Writing report...\")\n",
    "    input = f\"Original query: {query}\\n\\nSearch results: {search_results}\"\n",
    "    if previous_feedback:\n",
    "        input += f\"\\n\\nPrevious attempt was rejected. Feedback: {previous_feedback}\\nPlease improve the report based on this feedback.\"\n",
    "    result = await Runner.run(writer_agent, input)\n",
    "    print(\"Report complete\")\n",
    "    return result.final_output\n",
    "\n",
    "async def evaluate_report(query: str, report: str):\n",
    "    print(\"Evaluating report quality...\")\n",
    "    input = f\"Query: {query}\\n\\nReport:\\n{report}\"\n",
    "    result = await Runner.run(evaluator_agent, input)\n",
    "    evaluation = result.final_output\n",
    "    if evaluation.is_acceptable:\n",
    "        print(\"Report passed evaluation\")\n",
    "    else:\n",
    "        print(f\"Report needs improvement: {evaluation.feedback}\")\n",
    "    return evaluation\n",
    "\n",
    "async def send_report(report: ReportData):\n",
    "    print(\"Sending email...\")\n",
    "    result = await Runner.run(email_agent, report.markdown_report)\n",
    "    print(\"Email sent\")\n",
    "    return report"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63e7e366",
   "metadata": {},
   "source": [
    "## Main Research Agents Ochastration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00db6caa",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def enhanced_deep_research(query: str, user_clarifications: str = \"\"):\n",
    "    with trace(\"Enhanced Deep Research\"):\n",
    "        clarification = await clarify_query(query)\n",
    "        \n",
    "        if clarification.needs_clarification and not user_clarifications:\n",
    "            print(\"\\nQuery needs clarification. Please answer these questions:\")\n",
    "            for i, question in enumerate(clarification.questions, 1):\n",
    "                print(f\"{i}. {question}\")\n",
    "            print(\"\\nRerun this function with your clarifications as the second parameter.\")\n",
    "            return\n",
    "        \n",
    "        final_query = query\n",
    "        if user_clarifications:\n",
    "            final_query = f\"{query}\\n\\nAdditional context: {user_clarifications}\"\n",
    "        \n",
    "        search_plan = await plan_searches(final_query)\n",
    "        search_results = await perform_searches(search_plan)\n",
    "        \n",
    "        max_attempts = 2\n",
    "        for attempt in range(max_attempts):\n",
    "            if attempt == 0:\n",
    "                report = await write_report(final_query, search_results)\n",
    "            else:\n",
    "                report = await write_report(final_query, search_results, evaluation.feedback)\n",
    "            \n",
    "            evaluation = await evaluate_report(final_query, report.markdown_report)\n",
    "            \n",
    "            if evaluation.is_acceptable:\n",
    "                break\n",
    "            \n",
    "            if attempt == max_attempts - 1:\n",
    "                print(\"Max attempts reached. Sending report anyway.\")\n",
    "        \n",
    "        await send_report(report)\n",
    "        \n",
    "        display(Markdown(report.markdown_report))\n",
    "        print(\"\\nResearch complete!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63cd323e",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"AI Agent frameworks\"\n",
    "clarifications = \"Focus on production-ready frameworks, comparing OpenAI Agents SDK, LangGraph, and CrewAI. Looking for enterprise use cases in 2025.\"\n",
    "\n",
    "await enhanced_deep_research(query, clarifications)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
